//  Copyright 2011 Charlie Garrett
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//  http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//
//  Term.cpp
//

#include <sstream>

#include "assert.h"

#include "CodeGenerator.h"
#include "Instruction.h"
#include "Interpreter.h"
#include "Interpreter.h"
#include "Term.h"
#include "TermVisitor.h"

namespace logic {
    using namespace std;

    bool Functor::Equals(const Functor &other) const {
        return arity == other.arity && name == other.name;
    }
    
    bool Functor::operator<(const Functor &other) const {
        if (name < other.name) {
            return true;
        } else if (name > other.name) {
            return false;
        } else {
            return arity < other.arity;
        }
    }
    
    string Functor::ToString() const {
        ostringstream os;
        os << name << "/" << arity;
        return os.str();
    }
    
    Functor Functor::Undefined("", 0);
    Functor Functor::ListCons("cons", 2);
    
    void Term::GenerateProgramInstructions(CodeGenerator *codegen, deque<Term*> *worklist,
                                           bool as_argument) {
        assert(false);
    }
    
    void Term::GenerateQueryInstructions(CodeGenerator *codegen, bool as_argument, bool link_variables) {
        assert(false);
    }

    bool Constant::IsConstantNamed(const string &name) const {
        return value.kind == Value::STRING && (*value.data.str) == name;
    }
    
    void Constant::GetVariables(set<Variable*> *vars) {}
    
    void Constant::GenerateProgramInstructions(CodeGenerator *codegen, deque<Term*> *worklist,
                                               bool as_argument) {
        Register dest_reg;
        if (as_argument) {
            dest_reg = codegen->GetRegister(Register::ARGUMENT, this);
            assert(dest_reg.kind != Register::UNALLOCATED);
        } else {
            dest_reg = codegen->GetRegister(Register::TEMPORARY, this);
            if (dest_reg.kind == Register::UNALLOCATED) {
                dest_reg = codegen->AllocateRegister(Register::TEMPORARY, this);
            }
        }
        
        codegen->EmitInstruction(new GetConstant(value, dest_reg), this);
    }
    
    void Constant::GenerateQueryInstructions(CodeGenerator *codegen, bool as_argument, bool link_variables) {
        Register dest_reg;
        if (as_argument) {
            dest_reg = codegen->GetRegister(Register::ARGUMENT, this);
            assert(dest_reg.kind != Register::UNALLOCATED);
        } else {
            dest_reg = codegen->GetRegister(Register::TEMPORARY, this);
            if (dest_reg.kind == Register::UNALLOCATED) {
                dest_reg = codegen->AllocateRegister(Register::TEMPORARY, this);
            }
        }
        
        codegen->EmitInstruction(new PutConstant(value, dest_reg), this);
    }
    
    void Constant::Accept(TermVisitor* visitor) {
        visitor->Visit(this);
    }
    
    string Constant::ToString() const {
        return value.ToString();
    }
    
    void Variable::GetVariables(set<Variable*> *vars) {
        vars->insert(this);
    }
    
    void Variable::Accept(TermVisitor* visitor) {
        visitor->Visit(this);
    }
    
    void Variable::LinkVariable(CodeGenerator *codegen) {
        Register reg = codegen->GetRegister(Register::TEMPORARY, this);
        assert(reg.kind != Register::UNALLOCATED);
        codegen->EmitInstruction(new LinkSourceVariable(this, reg), this);
    }
    
    void Variable::GenerateProgramInstructions(CodeGenerator *codegen, deque<Term*> *worklist,
                                               bool as_argument) {
        // Code is only generated for Variables when they appear as top-level predicate
        // arguments.
        // If the Variable is allocated a permanent register, then use it.  Otherwise,
        // allocate a temporary register.
        Register normal_reg = codegen->GetRegister(Register::PERMANENT, this);
        if (normal_reg.kind == Register::UNALLOCATED) {
            normal_reg = codegen->GetRegister(Register::TEMPORARY, this);
            if (normal_reg.kind == Register::UNALLOCATED) {
                normal_reg = codegen->AllocateRegister(Register::TEMPORARY, this);
            }
        }

        if (as_argument) {
            // As arguments, Variables need both a normal and an argument register.
            // The argument register should already have been allocated by the
            // containing predicate.
            Register arg_reg = codegen->GetRegister(Register::ARGUMENT, this);
            assert(arg_reg.kind != Register::UNALLOCATED);

            if (codegen->IsRegisterUsed(normal_reg)) {
                codegen->EmitInstruction(new GetValue(normal_reg, arg_reg), this);
            } else {
                codegen->EmitInstruction(new GetVariable(normal_reg, arg_reg), this);
            }
        }
    }

    void Variable::GenerateQueryInstructions(CodeGenerator *codegen, bool as_argument, bool link_variables) {
        // Code is only generated for Variables when they appear as top-level predicate
        // arguments.
        // If the Variable is allocated a permanent register, then use it.  Otherwise,
        // allocate a temporary register.
        Register normal_reg = codegen->GetRegister(Register::PERMANENT, this);
        if (normal_reg.kind == Register::UNALLOCATED) {
            normal_reg = codegen->GetRegister(Register::TEMPORARY, this);
            if (normal_reg.kind == Register::UNALLOCATED) {
                normal_reg = codegen->AllocateRegister(Register::TEMPORARY, this);
            }
        }
        
        if (as_argument) {
            // As arguments, Variables need both a normal and an argument register.
            // The argument register should already have been allocated by the
            // containing predicate.
            Register arg_reg = codegen->GetRegister(Register::ARGUMENT, this);
            assert(arg_reg.kind != Register::UNALLOCATED);
            
            if (codegen->IsRegisterUsed(normal_reg)) {
                codegen->EmitInstruction(new PutValue(normal_reg, arg_reg), this);
            } else {
                codegen->EmitInstruction(new PutVariable(normal_reg, arg_reg), this);
                if (link_variables) {
                    LinkVariable(codegen);
                }
            }
        }
    }

    string Variable::ToString() const {
        return name;
    }
    
    // Add subterms to the worklist.
    void Structure::AddSubterms(deque<Term*> *worklist) {
        for (vector<Term*>::iterator sub = subterms.begin(); sub != subterms.end(); ++sub) {
            if (!(*sub)->IsConstant()) {
                worklist->push_back(*sub);
            }
        }        
    }
    
    Register Structure::DestReg(CodeGenerator *codegen, bool as_argument) {
        Register dest_reg;
        if (as_argument) {
            dest_reg = codegen->GetRegister(Register::ARGUMENT, this);
            assert(dest_reg.kind != Register::UNALLOCATED);
        } else {
            dest_reg = codegen->GetRegister(Register::TEMPORARY, this);
            if (dest_reg.kind == Register::UNALLOCATED) {
                dest_reg = codegen->AllocateRegister(Register::TEMPORARY, this);
            }
        }
        return dest_reg;
    }
    
    // Unify with or set references to the registers holding subterms.
    void Structure::UnifySubterms(CodeGenerator *codegen) {
        for (vector<Term*>::iterator sub = subterms.begin(); sub != subterms.end(); ++sub) {
            if ((*sub)->IsConstant()) {
                codegen->EmitInstruction(new UnifyConstant(*(*sub)->GetConstantValue()));
            } else {
                Register sub_reg = codegen->GetRegister(Register::TEMPORARY, *sub);
                if (sub_reg.kind == Register::UNALLOCATED) {
                    sub_reg = codegen->AllocateRegister(Register::TEMPORARY, *sub);
                }
                
                if (codegen->IsRegisterUsed(sub_reg)) {
                    codegen->EmitInstruction(new UnifyValue(sub_reg), *sub);
                } else {
                    codegen->EmitInstruction(new UnifyVariable(sub_reg), *sub);
                }
            }
        }
    }
    
    void Structure::GenerateSubtermQueryInstructions(CodeGenerator *codegen,
                                                     bool link_variables) {
        for (vector<Term*>::iterator sub = subterms.begin(); sub != subterms.end(); ++sub) {
            if (!(*sub)->IsConstant()) {
                (*sub)->GenerateQueryInstructions(codegen, false, link_variables);
            }
        }
    }
    
    void Structure::SetSubterms(CodeGenerator *codegen, bool link_variables) {
        // Set references to the registers holding subterms.
        for (vector<Term*>::iterator sub = subterms.begin(); sub != subterms.end(); ++sub) {
            if ((*sub)->IsConstant()) {
                codegen->EmitInstruction(new SetConstant(*(*sub)->GetConstantValue()));
            } else {
                Register sub_reg = codegen->GetRegister(Register::TEMPORARY, *sub);
                assert(sub_reg.kind != Register::UNALLOCATED);
                
                if (codegen->IsRegisterUsed(sub_reg)) {
                    codegen->EmitInstruction(new SetValue(sub_reg), *sub);
                } else {
                    codegen->EmitInstruction(new SetVariable(sub_reg), *sub);
                    if (link_variables) {
                        (*sub)->LinkVariable(codegen);
                    }
                }
            }
        }
    }
    
    void Structure::GetVariables(set<Variable*> *vars) {
        for (vector<Term*>::const_iterator sub = subterms.begin(); sub != subterms.end(); ++sub) {
            (*sub)->GetVariables(vars);
        }
    }
    
    void Structure::Accept(TermVisitor* visitor) {
        visitor->Visit(this);
    }
    
    void Structure::AddChildren(TermVisitor *visitor, bool reversed) {
        if (reversed) {
            for (vector<Term*>::reverse_iterator sub = subterms.rbegin(); sub != subterms.rend(); ++sub) {
                visitor->AddTerm(*sub);
            }            
        } else {
            for (vector<Term*>::iterator sub = subterms.begin(); sub != subterms.end(); ++sub) {
                visitor->AddTerm(*sub);
            }
        } 
    }
    
    void Structure::GenerateProgramInstructions(CodeGenerator *codegen,
                                                deque<Term*> *worklist,
                                                bool as_argument) {
        AddSubterms(worklist);
        
        // Unify with structure on the heap.
        codegen->EmitInstruction(new GetStructure(functor, DestReg(codegen, as_argument)),
                                 this);
        
        UnifySubterms(codegen);
    }
    
    void Structure::GenerateQueryInstructions(CodeGenerator *codegen,
                                              bool as_argument,
                                              bool link_variables) {
        GenerateSubtermQueryInstructions(codegen, link_variables);
        
        // Build a structure on the heap.
        codegen->EmitInstruction(new PutStructure(functor, DestReg(codegen, as_argument)),
                                 this);
        
        SetSubterms(codegen, link_variables);
    }
    
    string Structure::ToString() const {
        ostringstream os;
        os << functor.GetName();
        
        if (!subterms.empty()) {
            os << "(";
            
            bool needs_separator = false;
            for (vector<Term*>::const_iterator it = subterms.begin(); it != subterms.end(); ++it) {
                if (needs_separator) {
                    os << ", ";
                }
                
                os << (*it)->ToString();
                needs_separator = true;
            }
            os << ")";
        }
        return os.str();
    }
    
    void ListConsStructure::GenerateProgramInstructions(CodeGenerator *codegen,
                                                        deque<Term*> *worklist,
                                                        bool as_argument) {
        AddSubterms(worklist);

        // Unify with list on the heap.
        codegen->EmitInstruction(new GetList(DestReg(codegen, as_argument)), this);
        
        UnifySubterms(codegen);
    }
    
    void ListConsStructure::GenerateQueryInstructions(CodeGenerator *codegen,
                                                      bool as_argument,
                                                      bool link_variables) {
        GenerateSubtermQueryInstructions(codegen, link_variables);
        
        // Build a structure on the heap.
        codegen->EmitInstruction(new PutList(DestReg(codegen, as_argument)), this);
        
        SetSubterms(codegen, link_variables);
    }

    string ListConsStructure::ToString() const {
        ostringstream os;
        
        // Handle builtin functors specially.
        const Term *list = this;
        os << "[";
        while (list->IsStructureWithFunctor(Functor::ListCons)) {
            const Structure *list_struct = static_cast<const Structure*>(list);
            assert(list_struct->NumSubterms() == 2);
            // As long at the tail is a ListCons, iterate.
            os << list_struct->GetSubterm(0)->ToString();
            list = list_struct->GetSubterm(1);
            if (list->IsStructureWithFunctor(Functor::ListCons)) {
                os << ", ";
            } else if (!list->IsConstantNamed("[]")) {
                os << "|";
            }
        }
        // Skip printing empty lists in the tail.
        if (!list->IsConstantNamed("[]")) {
            os << list->ToString();
        }
        os << "]";
        return os.str();        
    }
    
    void Predicate::AllocateArgRegisters(CodeGenerator *codegen) {
        for (vector<Term*>::iterator sub = subterms.begin(); sub != subterms.end(); ++sub) {
            codegen->AllocateRegister(Register::ARGUMENT, *sub);
        }
    }
    
    void Predicate::Accept(TermVisitor* visitor) {
        visitor->Visit(this);
    }
    
    void Predicate::AddChildren(TermVisitor *visitor, bool reversed) {
        if (reversed) {
            for (vector<Term*>::reverse_iterator sub = subterms.rbegin(); sub != subterms.rend(); ++sub) {
                visitor->AddTerm(*sub);
            }            
        } else {
            for (vector<Term*>::iterator sub = subterms.begin(); sub != subterms.end(); ++sub) {
                visitor->AddTerm(*sub);
            }
        } 
    }
    
    void Predicate::TranslateQuery(CodeGenerator *codegen, bool link_variables) {
        AllocateArgRegisters(codegen);
        
        GenerateQueryInstructions(codegen, false, link_variables);
    }
    
    void Predicate::TranslateProgram(CodeGenerator *codegen) {
        AllocateArgRegisters(codegen);
        deque<Term*> worklist;
        GenerateProgramInstructions(codegen, &worklist);
    }
    
    void Predicate::GenerateProgramInstructions(CodeGenerator *codegen, deque<Term*> *worklist,
                                                bool as_argument) {
        // In program mode, we generate code for the arguments first, then for subterms in
        // breadth first order, using an auxiliary worklist.
        assert(worklist->size() == 0);
        for (vector<Term*>::iterator sub = subterms.begin(); sub != subterms.end(); ++sub) {
            (*sub)->GenerateProgramInstructions(codegen, worklist, true);
        }
        
        // Keep generating until the worklist is empty.
        while (worklist->size() > 0) {
            Term *sub = worklist->front();
            worklist->pop_front();
            sub->GenerateProgramInstructions(codegen, worklist);
        }
        
        // No register is needed for the emitted instruction.
        // Unify with arguments in registers and proceed.
        // codegen->EmitInstruction(new Proceed(), this);
        // codegen->EmitInstruction(new Halt(), this);
    }

    void Predicate::GenerateQueryInstructions(CodeGenerator *codegen, bool as_argument, bool link_variables) {
        for (vector<Term*>::iterator sub = subterms.begin(); sub != subterms.end(); ++sub) {
            (*sub)->GenerateQueryInstructions(codegen, true, link_variables);
        }
        // No register is needed for the emitted instruction.
        // Place arguments in registers and call the predicate.
        codegen->EmitInstruction(new Call(functor), this);
        // codegen->EmitInstruction(new Halt(), this);
    }
    
    void Clause::AllocatePermanentRegs(CodeGenerator *codegen) {
        // Variables that occur in multiple body atoms require permanent (caller-saved) registers.
        // The head of a clause is considered part of the first body atom for this calculation
        // because there is no call between them.
        
        if (body.size() > 1) {
            set<Variable*> prev_vars;
            set<Variable*> cur_vars;
        
            head->GetVariables(&prev_vars);
            body.at(0)->GetVariables(&prev_vars);
            
            for (int i = 1; i < body.size(); i++) {
                cur_vars.clear();
                body.at(i)->GetVariables(&cur_vars);
                
                // See which variables occur in this body atom and in some previous one.
                for (set<Variable*>::iterator cv = cur_vars.begin(); cv != cur_vars.end(); ++cv) {
                    if (prev_vars.find(*cv) != prev_vars.end()) {
                        Register reg = codegen->GetRegister(Register::PERMANENT, *cv);
                        if (reg.kind == Register::UNALLOCATED) {
                            reg = codegen->AllocateRegister(Register::PERMANENT, *cv);
                        }
                    }
                }

                for (set<Variable*>::iterator cv = cur_vars.begin(); cv != cur_vars.end(); ++cv) {
                    prev_vars.insert(*cv);
                }
            }
        }
    }
    
    void Clause::Translate(CodeGenerator *codegen, size_t total_clauses) {
        codegen->ClearRegisters();
        
        AllocatePermanentRegs(codegen);
        
        codegen->LabelInstruction(*head->GetFunctor(), clause_num);
        if (total_clauses > 1) {
            if (clause_num == 0) {
                codegen->EmitInstruction(new TryMeElse(*head->GetFunctor()));
            } else if (clause_num < total_clauses - 1) {
                codegen->EmitInstruction(new RetryMeElse(*head->GetFunctor(), clause_num));
            } else {
                codegen->EmitInstruction(new TrustMe());
            }
        }
        
        codegen->EmitInstruction(new Allocate(codegen->NumPermanentVars()));
        head->TranslateProgram(codegen);
        codegen->ClearArgRegisters();
        
        // Local registers are preserved between the head and the first body atom.
        for (vector<Predicate*>::iterator b = body.begin(); b != body.end(); ++b) {
            (*b)->TranslateQuery(codegen);
            // Argument and temporary registers are trashed by any call made by a body atom.
            codegen->ClearArgRegisters();
            codegen->ClearTempRegisters();
        }
        codegen->EmitInstruction(new Deallocate());
    }
    
    string Clause::ToString() const {
        ostringstream os;
        os << head->ToString();
        
        if (body.size() > 0) {
            os << " :- ";
        }
        
        bool needs_separator = false;
        for (vector<Predicate*>::const_iterator b = body.begin(); b != body.end(); ++b) {
            if (needs_separator) {
                os << ", ";
            }
            
            os << (*b)->ToString();
            needs_separator = true;
        }
        os << ".";
        
        return os.str();
    }
    
    void Rule::AddClause(Clause *clause) {
        clause->SetClauseNum(clauses.size());
        clauses.push_back(clause);
    }

    void Rule::Translate(CodeGenerator *codegen) {
        for (vector<Clause*>::iterator c = clauses.begin(); c < clauses.end(); ++c) {
            (*c)->Translate(codegen, clauses.size());
        }
    }
    
    string Rule::ToString() const {
        ostringstream os;
        for (vector<Clause*>::const_iterator c = clauses.begin(); c != clauses.end(); ++c) {
            os << (*c)->ToString() << endl;
        }
        return os.str();
    }

    void Program::AddClause(Clause *clause) {
        // TODO: Index Rules by functor for faster access.
        const Functor *clause_functor = clause->GetHead()->GetFunctor();
        for (vector<Rule*>::iterator r = rules.begin(); r != rules.end(); ++r) {
            const Functor *rule_functor = (*r)->GetClause(0)->GetHead()->GetFunctor();
            if (rule_functor->Equals(*clause_functor)) {
                (*r)->AddClause(clause);
                return;
            }
        }
        
        rules.push_back(new Rule(vector<Clause*>(1, clause)));
    }
    
    void Program::AddQuery(Predicate *query) {
        queries.push_back(query);
    }
    
    void Program::Translate(CodeGenerator *codegen) {
        for (vector<Rule*>::iterator r = rules.begin(); r != rules.end(); ++r) {
            (*r)->Translate(codegen);
        }
    }
    
    void Program::RunQuery(Predicate *query, bool trace) {
        // Create an interpreter, generate code for the whole program and the query,
        // then run it.
        Interpreter *interpreter = new Interpreter();
        interpreter->InitState();
        CodeGenerator codegen(interpreter);
        codegen.InitState();
        
        Translate(&codegen);
        // Use a dummy entry point for the query.
        Instruction::Label query_label =
            codegen.LabelInstruction(Functor("QEP", 0xffffffff), 0);
        query->TranslateQuery(&codegen, true);
        codegen.ResolveLabels();
        
        if (trace) {
            interpreter->PrintCode();
            interpreter->PrintState();
        }
        size_t offset = 0;
        if (codegen.LabelOffset(query_label, &offset)) {
            interpreter->RunFrom(offset);
        } else {
            cout << "Undefined predicate " << query_label.first.ToString() << endl;
        }
        if (trace) {
            interpreter->PrintCode();
            interpreter->PrintState();
        }
    }

    void Program::RunQueries(bool trace) {
        // Create an interpreter, generate code for the program rules.
        // Then generate code for each program query, if any, and run it.
        Interpreter *interpreter = new Interpreter(trace);
        interpreter->InitState();
        CodeGenerator codegen(interpreter);
        codegen.InitState();
        
        Translate(&codegen);
        interpreter->PrintCode();
        
        unsigned query_num = 0;
        for (vector<Predicate*>::iterator query = queries.begin();
             query != queries.end(); ++query) {
            Instruction::Label query_label =
                codegen.LabelInstruction(Functor("QEP", 0xffffffff), query_num);
            (*query)->TranslateQuery(&codegen, true);
            codegen.ResolveLabels();
            
            if (trace) {
                interpreter->PrintCode();
                interpreter->PrintState();
            }
            size_t offset = 0;
            if (codegen.LabelOffset(query_label, &offset)) {
                interpreter->RunFrom(offset);
            } else {
                cout << "Undefined predicate " << query_label.first.ToString() << endl;
            }
            if (trace) {
                interpreter->PrintCode();
                interpreter->PrintState();
            }
        }
    }
    
    string Program::ToString() const {
        ostringstream os;
        for (vector<Rule*>::const_iterator r = rules.begin(); r != rules.end(); ++r) {
            os << (*r)->ToString() << endl;
        }
        return os.str();
    }
    
}  // namespace logic

